{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\conda\\envs\\tuning_env\\lib\\site-packages\\torch\\cuda\\__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.\n",
      "  import pynvml  # type: ignore[import]\n",
      "d:\\conda\\envs\\tuning_env\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# step1 导入相关包\n",
    "import os\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# step2 加载本地模型和tokenizer\n",
    "model_path = r\"D:\\ghh\\model\\small\\Qwen\\Qwen2___5-0___5B-Instruct\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(model_path)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 我们动手实现`model.generate`方法的运行过程。\n",
    "# 定义generate函数, 作用是，预测下一个词, 也就是下一个Token\n",
    "\n",
    "def generate_next_token(model, generate_ids, temperature=0.7, debug=False):\n",
    "  print(\"input_ids:\", generate_ids) if debug else None\n",
    "  logits = model.forward(generate_ids).logits\n",
    "  print(\"logits: \", logits) if debug else None\n",
    "  \n",
    "  if temperature > 0:\n",
    "      probs = torch.softmax(logits[:, -1] / temperature, dim = -1)\n",
    "      print(\"probs: \", len(probs[0]), probs) if debug else None\n",
    "      next_token = torch.multinomial(probs[-1], num_samples=1)  # 按照概率采样得到下一个token\n",
    "  else:\n",
    "      next_token = torch.argmax(logits[:, -1], dim=-1)\n",
    "  \n",
    "  print(\"next_id: \", next_token, \", token: \", tokenizer.decode(next_token)) if debug else None\n",
    "  return next_token.reshape(-1, 1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 输入一个prompt 生成回答的步骤\n",
    "1. 输入Text --> Tokenizer 分词 --> input_ids --> Embedding\n",
    "1. 调用生成下一个token 的函数，将输入的文本进行编码，分词，生成向量，预测下一个词\n",
    "2. 再解码，输出"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tuning_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5b5cfca19401015264b4741f08a5ab630bb44e69b46d09de08cfb7d2d4656fc3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
